from typing import List

import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split


def drop(train_df, test_df, bye: List):
    for i in [train_df, test_df]:
        for z in bye:
            del i[z]
    return train_df, test_df


def remove_spaces(s):
    return ''.join(str(s).split())


def one_hot_encode(train_df, test_df, columns: List):
    for column in columns:
        train_df[column] = train_df[column].apply(lambda x: str(x))
        test_df[column] = test_df[column].apply(lambda x: str(x))
        good_cols = [column + '_' + i for i in train_df[column].unique() if i in test_df[column].unique()]
        train_df = pd.concat((train_df, pd.get_dummies(train_df[column], prefix=column)[good_cols]), axis=1)
        test_df = pd.concat((test_df, pd.get_dummies(test_df[column], prefix=column)[good_cols]), axis=1)
        del train_df[column]
        del test_df[column]
    return train_df, test_df


df = pd.read_excel('./data/liverpea100.xlsx', engine='openpyxl')

# Shuffle the DataFrame
df_shuffled = df.sample(frac=1, random_state=1)  # Set random_state for reproducibility

# Split the shuffled data into training and testing datasets
train_df, test_df = train_test_split(df_shuffled, test_size=0.1, random_state=1)

train_df = train_df.fillna(0)
test_df = test_df.fillna(0)

train_df['初步判断中医证型'] = train_df['初步判断中医证型'].apply(lambda x: str(x).split('┋')[0])
test_df['初步判断中医证型'] = test_df['初步判断中医证型'].apply(lambda x: str(x).split('┋')[0])

train_df, test_df = drop(train_df, test_df,
                         ['编号', '病程', '是否首次确诊', '地域（哪个省）', '首发症状', 'A望诊其他', 'B闻诊其他',
                          'C问诊其他', 'D按诊其他', 'E舌诊其他', 'F脉诊其他', '其他中医证型'])

train_df, test_df = one_hot_encode(train_df, test_df,
                                   ['性别', '职业', '受教育程度', '临床分型', '家族史', '初步判断中医证型'])

rf = RandomForestClassifier(max_features='sqrt',
                            criterion='gini',
                            n_estimators=700,
                            min_samples_split=2,
                            min_samples_leaf=1,
                            oob_score=True,
                            random_state=1,
                            n_jobs=-1)

train_inp = train_df.iloc[:, 0:186]
train_oup = train_df.iloc[:, 186:]
rf.fit(train_inp, train_oup)
print("%.4f" % rf.oob_score_)

test_inp = test_df.iloc[:, 0:186]
prediction = rf.predict(test_inp)

prediction_train = rf.predict(train_inp)

print()
